package com.bo.day20241121;

import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
 * Prim算法生成最小树
 * 思路是：
 * 1. 从任意一个点开始，将这个点加入到最小树中，并标记为已访问
 * 点 -> 解锁边(寻   找最小的那个边)  ->  解锁点 -> 解锁边 -> ...
 *
 * @Author: gpb
 * @Date: 2024/11/27 15:52
 * @Description:
 */
public class Prim {
    // 边排序
    public static class EdgeComparator implements Comparator<Edge> {
        public int compare (Edge o1, Edge o2) {
            return o1.width - o2.width;
        }
    }

    public static Set<Edge> prim (Graph graph) {
        // 解锁的边进入小根堆
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());

        // 哪些点被解锁出来了
        HashSet<Node> nodeSet = new HashSet<>();

        Set<Edge> result = new HashSet<>();
        // 随便挑选一个点开始
        for (Node node : graph.nodes.values()) {
            if (!nodeSet.contains(node)) {
                // 将解锁的边加入小根堆
                for (Edge edge : node.edges) {
                    priorityQueue.add(edge);
                }
                while (!priorityQueue.isEmpty()) {
                    Edge edge = priorityQueue.poll();
                    Node toNode = edge.to;
                    // 如果这个点没有被解锁过，就加入到最小树中，并标记为已访问
                    if (!nodeSet.contains(toNode)) {
                        result.add(edge);
                        nodeSet.add(toNode);
                        // 解锁新的边
                        for (Edge newEdge : toNode.edges) {
                            priorityQueue.add(newEdge);
                        }
                    }
                }
            }
        }
        return result;
    }
}
